


#include "yara/tlsh.h"
#include "yara/tlsh_impl.h"
#include "yara/tlsh_util.h"



#include<malloc.h>


#include <string.h>
#include <errno.h>
unsigned int *a_bucket;
unsigned char slide_window[SLIDING_WND_SIZE];
unsigned int data_len_tlsh;

struct lsh_bin_struct {
    unsigned char checksum[TLSH_CHECKSUM_LEN];  // 1 to 3 bytes
    unsigned char Lvalue;                       // 1 byte
    union {
        unsigned char QB;
        struct{
            unsigned char Q1ratio : 4;
            unsigned char Q2ratio : 4;
        } QR;
    } Q;                                        // 1 bytes
    unsigned char tmp_code[CODE_SIZE];          // 32/64 bytes
} lsh_bin, lsh_bin2, lsh_bintest, lsh_bintest_tmp, lsh_bin_hex;

char *lsh_code;       // allocated when hash() function without buffer is called - 70/134 bytes or 74/138 bytes
int lsh_code_valid;  // true iff final() or fromTlshStr complete successfully

static void find_quartile(unsigned int *q1, unsigned int *q2, unsigned int *q3, const unsigned int * a_bucket);
static unsigned int partition(unsigned int * buf, unsigned int left, unsigned int right);

void update1(const  unsigned char* data, unsigned int len)
{
#define RNG_SIZE    	SLIDING_WND_SIZE
#define RNG_IDX(i)	((i+RNG_SIZE)%RNG_SIZE)

    int j = (int)(data_len_tlsh % RNG_SIZE);
    unsigned int fed_len = data_len_tlsh;

    if (a_bucket == NULL) {
        a_bucket = (unsigned int *)malloc(sizeof(unsigned int)*BUCKETS);
        memset(a_bucket, 0, sizeof(unsigned int)*BUCKETS);
    }

    for (unsigned int i = 0; i<len; i++, fed_len++, j = RNG_IDX(j + 1)) {
        slide_window[j] = data[i];

        if (fed_len >= 4) {
            //only calculate when input >= 5 bytes
            int j_1 = RNG_IDX(j - 1);
            int j_2 = RNG_IDX(j - 2);
            int j_3 = RNG_IDX(j - 3);
            int j_4 = RNG_IDX(j - 4);

            for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
                if (k == 0) {
                    lsh_bin.checksum[k] = b_mapping(0, slide_window[j], slide_window[j_1], lsh_bin.checksum[k]);
                }
                else {
                    // use calculated 1 byte checksums to expand the total checksum to 3 bytes
                    lsh_bin.checksum[k] = b_mapping(lsh_bin.checksum[k - 1], slide_window[j], slide_window[j_1], lsh_bin.checksum[k]);
                }
            }

            unsigned char r;
            r = b_mapping(2, slide_window[j], slide_window[j_1], slide_window[j_2]);
            a_bucket[r]++;
            r = b_mapping(3, slide_window[j], slide_window[j_1], slide_window[j_3]);
            a_bucket[r]++;
            r = b_mapping(5, slide_window[j], slide_window[j_2], slide_window[j_3]);
            a_bucket[r]++;
            r = b_mapping(7, slide_window[j], slide_window[j_2], slide_window[j_4]);
            a_bucket[r]++;
            r = b_mapping(11, slide_window[j], slide_window[j_1], slide_window[j_4]);
            a_bucket[r]++;
            r = b_mapping(13, slide_window[j], slide_window[j_3], slide_window[j_4]);
            a_bucket[r]++;

        }
    }
    data_len_tlsh += len;
}

/* to signal the class there is no more data to be added */
void final1()
{
    // incoming data must more than or equal to MIN_DATA_LENGTH bytes
    if (data_len_tlsh < MIN_DATA_LENGTH) {
        // lsh_code be empty
        //delete[] a_bucket; a_bucket = NULL;
        free(a_bucket);
        a_bucket = NULL;
        return;
    }

    unsigned int q1, q2, q3;
    find_quartile(&q1, &q2, &q3, a_bucket);

    // buckets must be more than 50% non-zero
    int nonzero = 0;
    for (unsigned int i = 0; i<CODE_SIZE; i++) {
        for (unsigned int j = 0; j<4; j++) {
            if (a_bucket[4 * i + j] > 0) {
                nonzero++;
            }
        }
    }
    if (nonzero <= 4 * CODE_SIZE / 2) {
        free(a_bucket);
        a_bucket = NULL;
        return;
    }

    for (unsigned int i = 0; i<CODE_SIZE; i++) {
        unsigned char h = 0;
        for (unsigned int j = 0; j<4; j++) {
            unsigned int k = a_bucket[4 * i + j];
            if (q3 < k) {
                h += 3 << (j * 2);  // leave the optimization j*2 = j<<1 or j*2 = j+j for compiler
            }
            else if (q2 < k) {
                h += 2 << (j * 2);
            }
            else if (q1 < k) {
                h += 1 << (j * 2);
            }
        }
        lsh_bin.tmp_code[i] = h;
    }

    //Done with a_bucket so deallocate
    //delete[] a_bucket; a_bucket = NULL;
    free(a_bucket);
    a_bucket = NULL;
    lsh_bin.Lvalue = l_capturing(data_len_tlsh);
    lsh_bin.Q.QR.Q1ratio = (unsigned int)((float)(q1 * 100) / (float)q3) % 16;
    lsh_bin.Q.QR.Q2ratio = (unsigned int)((float)(q2 * 100) / (float)q3) % 16;
    lsh_code_valid = 1;
}





void update1test(unsigned char* data, unsigned int len)
{
#define RNG_SIZE    	SLIDING_WND_SIZE
#define RNG_IDX(i)	((i+RNG_SIZE)%RNG_SIZE)

    int j = (int)(data_len_tlsh % RNG_SIZE);
    unsigned int fed_len = data_len_tlsh;

    if (a_bucket == NULL) {
        a_bucket = (unsigned int *)malloc(sizeof(unsigned int)*BUCKETS);
        memset(a_bucket, 0, sizeof(unsigned int)*BUCKETS);
    }

    for (unsigned int i = 0; i<len; i++, fed_len++, j = RNG_IDX(j + 1)) {
        slide_window[j] = data[i];

        if (fed_len >= 4) {
            //only calculate when input >= 5 bytes
            int j_1 = RNG_IDX(j - 1);
            int j_2 = RNG_IDX(j - 2);
            int j_3 = RNG_IDX(j - 3);
            int j_4 = RNG_IDX(j - 4);

            for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
                if (k == 0) {
                    lsh_bintest.checksum[k] = b_mapping(0, slide_window[j], slide_window[j_1], lsh_bintest.checksum[k]);
                }
                else {
                    // use calculated 1 byte checksums to expand the total checksum to 3 bytes
                    lsh_bintest.checksum[k] = b_mapping(lsh_bintest.checksum[k - 1], slide_window[j], slide_window[j_1], lsh_bintest.checksum[k]);
                }
            }

            unsigned char r;
            r = b_mapping(2, slide_window[j], slide_window[j_1], slide_window[j_2]);
            a_bucket[r]++;
            r = b_mapping(3, slide_window[j], slide_window[j_1], slide_window[j_3]);
            a_bucket[r]++;
            r = b_mapping(5, slide_window[j], slide_window[j_2], slide_window[j_3]);
            a_bucket[r]++;
            r = b_mapping(7, slide_window[j], slide_window[j_2], slide_window[j_4]);
            a_bucket[r]++;
            r = b_mapping(11, slide_window[j], slide_window[j_1], slide_window[j_4]);
            a_bucket[r]++;
            r = b_mapping(13, slide_window[j], slide_window[j_3], slide_window[j_4]);
            a_bucket[r]++;

        }
    }
    data_len_tlsh += len;
}

/* to signal the class there is no more data to be added */
void final1test()
{
    // incoming data must more than or equal to MIN_DATA_LENGTH bytes
    if (data_len_tlsh < MIN_DATA_LENGTH) {
        // lsh_code be empty
        //delete[] a_bucket; a_bucket = NULL;
        return;
    }

    unsigned int q1, q2, q3;
    find_quartile(&q1, &q2, &q3, a_bucket);

    // buckets must be more than 50% non-zero
    int nonzero = 0;
    for (unsigned int i = 0; i<CODE_SIZE; i++) {
        for (unsigned int j = 0; j<4; j++) {
            if (a_bucket[4 * i + j] > 0) {
                nonzero++;
            }
        }
    }
    if (nonzero <= 4 * CODE_SIZE / 2) {
        //delete[] a_bucket; a_bucket = NULL;
        return;
    }

    for (unsigned int i = 0; i<CODE_SIZE; i++) {
        unsigned char h = 0;
        for (unsigned int j = 0; j<4; j++) {
            unsigned int k = a_bucket[4 * i + j];
            if (q3 < k) {
                h += 3 << (j * 2);  // leave the optimization j*2 = j<<1 or j*2 = j+j for compiler
            }
            else if (q2 < k) {
                h += 2 << (j * 2);
            }
            else if (q1 < k) {
                h += 1 << (j * 2);
            }
        }
        lsh_bintest.tmp_code[i] = h;
    }

    //Done with a_bucket so deallocate
    //delete[] a_bucket; a_bucket = NULL;

    lsh_bintest.Lvalue = l_capturing(data_len_tlsh);
    lsh_bintest.Q.QR.Q1ratio = (unsigned int)((float)(q1 * 100) / (float)q3) % 16;
    lsh_bintest.Q.QR.Q2ratio = (unsigned int)((float)(q2 * 100) / (float)q3) % 16;
    lsh_code_valid = 1;
}

char* hash1(char *buffer, unsigned int bufSize)
{
    if (bufSize < TLSH_STRING_LEN + 1) {
        strncpy(buffer, "", bufSize);
        return buffer;
    }
    if (lsh_code_valid == 0) {
        strncpy(buffer, "", bufSize);
        return buffer;
    }


    //lsh_bin_struct lsh_bin11;

    for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
        lsh_bin2.checksum[k] = swap_byte(lsh_bin.checksum[k]);
    }
    lsh_bin2.Lvalue = swap_byte(lsh_bin.Lvalue);
    lsh_bin2.Q.QB = swap_byte(lsh_bin.Q.QB);
    for (int i = 0; i < CODE_SIZE; i++){
        lsh_bin2.tmp_code[i] = (lsh_bin.tmp_code[CODE_SIZE - 1 - i]);
    }

    to_hex((unsigned char*)&lsh_bin2, sizeof(lsh_bin2), buffer);

    return buffer;
}


void find_quartile(unsigned int *q1, unsigned int *q2, unsigned int *q3, const unsigned int * a_bucket)
{
    unsigned int bucket_copy[EFF_BUCKETS], short_cut_left[EFF_BUCKETS], short_cut_right[EFF_BUCKETS], spl = 0, spr = 0;
    unsigned int p1 = EFF_BUCKETS / 4 - 1;
    unsigned int p2 = EFF_BUCKETS / 2 - 1;
    unsigned int p3 = EFF_BUCKETS - EFF_BUCKETS / 4 - 1;
    unsigned int end = EFF_BUCKETS - 1;

    for (unsigned int i = 0; i <= end; i++) {
        bucket_copy[i] = a_bucket[i];
    }

    for (unsigned int l = 0, r = end;;) {
        unsigned int ret = partition(bucket_copy, l, r);
        if (ret > p2) {
            r = ret - 1;
            short_cut_right[spr] = ret;
            spr++;
        }
        else if (ret < p2){
            l = ret + 1;
            short_cut_left[spl] = ret;
            spl++;
        }
        else {
            *q2 = bucket_copy[p2];
            break;
        }
    }

    short_cut_left[spl] = p2 - 1;
    short_cut_right[spr] = p2 + 1;

    for (unsigned int i = 0, l = 0; i <= spl; i++) {
        unsigned int r = short_cut_left[i];
        if (r > p1) {
            for (;;) {
                unsigned int ret = partition(bucket_copy, l, r);
                if (ret > p1) {
                    r = ret - 1;
                }
                else if (ret < p1) {
                    l = ret + 1;
                }
                else {
                    *q1 = bucket_copy[p1];
                    break;
                }
            }
            break;
        }
        else if (r < p1) {
            l = r;
        }
        else {
            *q1 = bucket_copy[p1];
            break;
        }
    }

    for (unsigned int i = 0, r = end; i <= spr; i++) {
        unsigned int l = short_cut_right[i];
        if (l < p3) {
            for (;;) {
                unsigned int ret = partition(bucket_copy, l, r);
                if (ret > p3) {
                    r = ret - 1;
                }
                else if (ret < p3) {
                    l = ret + 1;
                }
                else {
                    *q3 = bucket_copy[p3];
                    break;
                }
            }
            break;
        }
        else if (l > p3) {
            r = l;
        }
        else {
            *q3 = bucket_copy[p3];
            break;
        }
    }

}


#define SWAP_UINT(x,y) do {\
    unsigned int int_tmp = (x);  \
    (x) = (y); \
    (y) = int_tmp; } while (0)


unsigned int partition(unsigned int * buf, unsigned int left, unsigned int right)
{
    if (left == right) {
        return left;
    }
    if (left + 1 == right) {
        if (buf[left] > buf[right]) {
            SWAP_UINT(buf[left], buf[right]);
        }
        return left;
    }

    unsigned int ret = left, pivot = (left + right) >> 1;

    unsigned int val = buf[pivot];

    buf[pivot] = buf[right];
    buf[right] = val;

    for (unsigned int i = left; i < right; i++) {
        if (buf[i] < val) {
            SWAP_UINT(buf[ret], buf[i]);
            ret++;
        }
    }
    buf[right] = buf[ret];
    buf[ret] = val;

    return ret;
}


char* hash1test(char *buffer, unsigned int bufSize)
{
    if (bufSize < TLSH_STRING_LEN + 1) {
        strncpy(buffer, "", bufSize);
        return buffer;
    }
    if (lsh_code_valid == 0) {
        strncpy(buffer, "", bufSize);
        return buffer;
    }


    //lsh_bin_struct lsh_bin11;

    for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
        lsh_bintest_tmp.checksum[k] = swap_byte(lsh_bintest.checksum[k]);
    }
    lsh_bintest_tmp.Lvalue = swap_byte(lsh_bintest.Lvalue);
    lsh_bintest_tmp.Q.QB = swap_byte(lsh_bintest.Q.QB);
    for (int i = 0; i < CODE_SIZE; i++){
        lsh_bintest_tmp.tmp_code[i] = (lsh_bintest.tmp_code[CODE_SIZE - 1 - i]);
    }

    to_hex((unsigned char*)&lsh_bintest_tmp, sizeof(lsh_bintest_tmp), buffer);

    return buffer;
}


int fromTlshStr_impl(const char* str)
{


    // Validate input string
    for (int i = 0; i < TLSH_STRING_LEN; i++)
    if (!(
        (str[i] >= '0' && str[i] <= '9') ||
        (str[i] >= 'A' && str[i] <= 'F') ||
        (str[i] >= 'a' && str[i] <= 'f')))
    {
        return 1;
    }

    //reset();


    from_hex(str, TLSH_STRING_LEN, (unsigned char*)&lsh_bin_hex);

    // Reconstruct checksum, Qrations & lvalue
    for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
        lsh_bintest.checksum[k] = swap_byte(lsh_bin_hex.checksum[k]);
    }
    lsh_bintest.Lvalue = swap_byte(lsh_bin_hex.Lvalue);
    lsh_bintest.Q.QB = swap_byte(lsh_bin_hex.Q.QB);
    for (int i = 0; i < CODE_SIZE; i++){
        lsh_bintest.tmp_code[i] = (lsh_bin_hex.tmp_code[CODE_SIZE - 1 - i]);
    }
    lsh_code_valid = 1;

    return 0;




}

int totalDiff(int len_diff)
{
    int diff = 0;

    if (len_diff) {
        int ldiff = mod_diff(lsh_bin.Lvalue, lsh_bintest.Lvalue, 256);
        if (ldiff == 0)
            diff = 0;
        else if (ldiff == 1)
            diff = 1;
        else
            diff += ldiff * 12;
    }

    int q1diff = mod_diff(lsh_bin.Q.QR.Q1ratio, lsh_bintest.Q.QR.Q1ratio, 16);
    if (q1diff <= 1)
        diff += q1diff;
    else
        diff += (q1diff - 1) * 12;

    int q2diff = mod_diff(lsh_bin.Q.QR.Q2ratio, lsh_bintest.Q.QR.Q2ratio, 16);
    if (q2diff <= 1)
        diff += q2diff;
    else
        diff += (q2diff - 1) * 12;

    for (int k = 0; k < TLSH_CHECKSUM_LEN; k++) {
        if (lsh_bin.checksum[k] != lsh_bintest.checksum[k]) {
            diff++;
            break;
        }
    }

    diff += h_distance(CODE_SIZE, lsh_bin.tmp_code, lsh_bintest.tmp_code);

    memset(&lsh_bin, 0, sizeof(lsh_bin));
    memset(&lsh_bintest, 0, sizeof(lsh_bintest));
    memset(&lsh_bin2, 0, sizeof(lsh_bin));
    memset(&lsh_bintest_tmp, 0, sizeof(lsh_bintest));
    memset(&lsh_bin_hex, 0, sizeof(lsh_bin));
    data_len_tlsh = 0;

    return (diff);
}
